import sys
from PyQt5.QtWidgets import QDialog, QMainWindow, QApplication, QMessageBox, QFileDialog, QInputDialog
from PyQt5 import QtGui
from mywin import Ui_MainWindow
import os
import os.path as osp
import time
from personTransform import personTransform1
from PIL import Image, ImageQt


class MyMainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        super(MyMainWindow, self).__init__(parent)
        self.setupUi(self)
        self.boundFun()
        self.path_imgs = []
        self.index = 0
        self.time_slot = 0.5
        self.stop_token = False
        self.transform = None

    def boundFun(self):
        self.actionOpenImgDir.triggered.connect(self.openImgDir)

        self.actionCUHK03_detected.triggered.connect(
            self.selectCUHK03_detected)
        self.actionCUHK03_labeled.triggered.connect(self.selectCUHK03_labeled)

        self.actionmarket1501_query.triggered.connect(
            self.selectMarket1501_query)
        self.actionmarket1501_bbox_train.triggered.connect(
            self.selectMarket1501_train)
        self.actionmarket1501_bbox_test.triggered.connect(
            self.selectMarket1501_test)
        self.actionmarket1501_gt_bbox.triggered.connect(
            self.selectMarket1501_gt_bbox)

        self.actionDukeMTMC_train.triggered.connect(self.selectDukeMTMC_train)
        self.actionDukeMTMC_query.triggered.connect(self.selectDukeMTMC_query)
        self.actionDukeMTMC_test.triggered.connect(self.selectDukeMTMC_test)

        self.actionOccluded_DukeMTMC.triggered.connect(
            self.selectOccluded_DukeMTMC)

        self.button_jump.clicked.connect(self.jump_to_pic)
        self.button_next.clicked.connect(self.next_pic)
        self.button_previous.clicked.connect(self.previous_pic)

        self.horizontalSlider.valueChanged.connect(self.update_Speed)
        self.button_autoplay.clicked.connect(self.autoplay)
        self.button_stopplay.clicked.connect(self.stopplay)

        self.button_manual_a.clicked.connect(self.set_transform_none)
        self.button_random_a.clicked.connect(self.set_transform1)

    def showMessageBox(self, title: str, content: str):
        ret = QMessageBox.information(
            self, title, content, QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
        return ret

    def set_transform_none(self):
        self.transform = None
        self.label_triplet_sampling.setText('transform: None')
        try:
            self.display_label_1(self.index)
        except:
            pass

    def set_transform1(self):
        self.transform = personTransform1
        self.label_triplet_sampling.setText('transform: Type 1')
        try:
            self.display_label_1(self.index)
        except:
            pass

    def display_label_1(self, the_index):
        path_img = self.path_imgs[the_index]

        image = Image.open(path_img)
        if self.transform is not None:
            image = self.transform(image)
        qimage = ImageQt.ImageQt(image)
        img = QtGui.QPixmap.fromImage(qimage)

        # img = QtGui.QPixmap(path_img).scaled(self.label.width(), self.label.height())

        self.label.setPixmap(img)
        self.label_index.setText(f"序号 {the_index}")
        self.icon.setText(
            f"{osp.basename(osp.dirname(path_img))}/{osp.basename(path_img)}")

    def jump_to_pic(self):
        num, ok = QInputDialog.getInt(self, 'Integer input dialog', '输入数字')
        if ok and num:
            self.index = num
            try:
                self.display_label_1(self.index)
            except:
                pass

    def next_pic(self):
        leng = len(self.path_imgs)
        if leng == 0:
            return
        if self.index < leng-1:
            self.index += 1
        self.display_label_1(self.index)

    def previous_pic(self):
        if len(self.path_imgs) == 0:
            return
        if self.index > 0:
            self.index -= 1
        self.display_label_1(self.index)

    def update_Speed(self):
        self.time_slot = self.horizontalSlider.value() * 0.1

    def stopplay(self):
        self.stop_token = True

    def autoplay(self):
        self.stop_token = False
        leng = len(self.path_imgs)
        if leng == 0:
            return
        for i in range(self.index, leng):
            self.index = i
            QApplication.processEvents()
            self.display_label_1(self.index)
            QApplication.processEvents()
            time.sleep(self.time_slot)
            if self.stop_token:
                break

    def read_imgDir(self, path_imgDir):
        self.index = 0
        path_imgDir = osp.abspath(path_imgDir)
        imgnames = [name for name in os.listdir(path_imgDir) if name.split(
            '.')[-1] in ['jpg', 'jpeg', 'png', 'bmp', 'tif']]
        path_imgs = []
        for imgname in imgnames:
            path_img = osp.join(path_imgDir, imgname)
            path_imgs.append(path_img)
        return path_imgs

    def openImgDir(self):
        self.icon_2.setText("手动选择数据集")
        self.path_imageset = QFileDialog.getExistingDirectory(
            self, "选择图片集文件夹", ".")
        print("选择图片文件夹路径:\n   ", osp.abspath(self.path_imageset))
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectCUHK03_detected(self):
        self.icon_2.setText("CUHK03_detected 数据集")
        self.path_imageset = r'D:\reid\DATA4REID\cuhk03\images_detected'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectCUHK03_labeled(self):
        self.icon_2.setText("CUHK03_labeled 数据集")
        self.path_imageset = r'D:\reid\DATA4REID\cuhk03\images_labeled'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectMarket1501_query(self):
        self.icon_2.setText("market1501 数据集 query")
        self.path_imageset = r'D:\reid\DATA4REID\Market-1501-v15.09.15\query'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectMarket1501_train(self):
        self.icon_2.setText("market1501 数据集 bbox_train")
        self.path_imageset = r'D:\reid\DATA4REID\Market-1501-v15.09.15\bounding_box_train'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectMarket1501_test(self):
        self.icon_2.setText("market1501 数据集 bbox_test")
        self.path_imageset = r'D:\reid\DATA4REID\Market-1501-v15.09.15\bounding_box_test'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectMarket1501_gt_bbox(self):
        self.icon_2.setText("market1501 数据集 gt_bbox")
        self.path_imageset = r'D:\reid\DATA4REID\Market-1501-v15.09.15\gt_bbox'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectDukeMTMC_train(self):
        self.icon_2.setText("DukeMTMC-reID 数据集 bbox_train")
        self.path_imageset = r'D:\reid\DATA4REID\DukeMTMC-reID\bounding_box_train'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectDukeMTMC_query(self):
        self.icon_2.setText("DukeMTMC-reID 数据集 query")
        self.path_imageset = r'D:\reid\DATA4REID\DukeMTMC-reID\query'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectDukeMTMC_test(self):
        self.icon_2.setText("DukeMTMC-reID 数据集 bbox_test")
        self.path_imageset = r'D:\reid\DATA4REID\DukeMTMC-reID\bounding_box_test'
        self.path_imgs = self.read_imgDir(self.path_imageset)
        self.display_label_1(self.index)

    def selectOccluded_DukeMTMC(self):
        self.showMessageBox("Occluded_DukeMTMC", "...")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    ui = MyMainWindow()
    ui.show()
    sys.exit(app.exec_())
